/*
 * Copyright (c) 2025 [ByteDance Ltd. and/or its affiliates.] and/or its affiliates
 * SPDX-License-Identifier: LGPL v2.1+
*/
#include "FFmpegXORIODecryptor.h"
#include "FFmpegIOCryptorUtils.h"
#include "FFErrorCodeDef.h"

int FFmpegXORIODecryptor::prepare(const char *url, AVFormatContext *formatContext, const AVDictionary *options, const std::map<std::string, FFAny> *params) {
    auto ret = FFmpegBaseIOCryptor::prepare(url, formatContext, options, params);
    if (ret < 0) return ret;

    auto oldPos = avio_tell(m_pOriginalIOCtx);

    uint32_t tagSize = 0;
    ret = findBdveTag(&tagSize);
    if (ret < 0) {
        m_shouldDecrypt = true;
        return ret;
    }

    m_bdveTagSize = tagSize;
    m_maxReadPos = m_fileSize - m_bdveTagSize;

    ret = parseBdveTag(tagSize);
    if (ret < 0) {
        m_shouldDecrypt = true;
        return ret;
    }

    avio_flush(m_pOriginalIOCtx);
    auto pos = avio_seek(m_pOriginalIOCtx, oldPos, SEEK_SET);
    if (pos < 0) {
        m_shouldDecrypt = true;
        return static_cast<int>(pos);
    }

    return 0;
}

int FFmpegXORIODecryptor::findBdveTag(uint32_t *tagSize) {
    auto size = avio_size(m_pOriginalIOCtx);
    if (size < 0) {
        av_log(NULL, AV_LOG_ERROR, "avio_size failed with error: %" PRId64 "", size);
        return static_cast<int>(size);
    }

    m_fileSize = size;

    if (size < 8) {
        av_log(NULL, AV_LOG_ERROR, "Invalid tag size");
        return TE_ERR_DECRYPT_INVALID_TAG_SIZE;
    }

    auto pos = avio_seek(m_pOriginalIOCtx, size - 8, SEEK_SET);
    if (pos < 0) {
        av_log(NULL, AV_LOG_ERROR, "avio_seek failed with error: %" PRId64 "", pos);
        return static_cast<int>(pos);
    }

    auto tag = avio_rb32(m_pOriginalIOCtx);
    if (tag != MKBETAG('s', 'i', 'z', 'e')) {
        av_log(NULL, AV_LOG_ERROR, "Invalid tag(need size tag)");
        return TE_ERR_DECRYPT_INVALID_TAG;
    }

    auto bdveTagSize = avio_rb32(m_pOriginalIOCtx);
    if (size < bdveTagSize) {
        av_log(NULL, AV_LOG_ERROR, "Invalid tag size");
        return TE_ERR_DECRYPT_INVALID_TAG_SIZE;
    }

    pos = avio_seek(m_pOriginalIOCtx, size - bdveTagSize + 4, SEEK_SET);
    if (pos < 0) {
        av_log(NULL, AV_LOG_ERROR, "avio_seek failed with error: %" PRId64 "", pos);
        return static_cast<int>(pos);
    }

    tag = avio_rb32(m_pOriginalIOCtx);
    if (tag != MKBETAG('b', 'd', 'v', 'e')) {
        av_log(NULL, AV_LOG_ERROR, "Invalid tag(need bdve tag)");
        return TE_ERR_DECRYPT_INVALID_TAG;
    }

    *tagSize = bdveTagSize;
    return 0;
}

int FFmpegXORIODecryptor::parseBdveTag(const uint32_t tagSize) {
    uint32_t crptTagSize = 0;
    if (FFmpegIOCryptorUtils::findTag(m_pOriginalIOCtx, MKBETAG('c', 'r', 'p', 't'), tagSize - 8, &crptTagSize)) {
        auto cryptType = avio_rb32(m_pOriginalIOCtx);
        auto cryptVersion = avio_rb32(m_pOriginalIOCtx);
        if (cryptType == TECryptorType::TECryptorType_XOR) {
            if (CRYPT_VERSION_2 == cryptVersion) {
                //TODO:
            } else if (CRYPT_VERSION_3 == cryptVersion) {
                uint8_t dst[32];
                encryptParamSha256(dst);

                uint8_t sha256[32];
                auto len = avio_read(m_pOriginalIOCtx, sha256, 32);
                if (len != 32) {
                    av_log(NULL, AV_LOG_ERROR, "Read sha256 data failed!");
                    return TE_ERR_DECRYPT_READ_SHA256_FAILED;
                }

                if (memcmp(dst, sha256, 32) != 0) {
                    av_log(NULL, AV_LOG_ERROR, "SHA-256 data not match!");
                    return TE_ERR_DECRYPT_SHA256_NOT_MATCH;
                }
            }

            return 0;
        } else {
            av_log(NULL, AV_LOG_ERROR, "cryptor type [%" PRIu32 "] not supported", cryptType);
            return TE_ERR_DECRYPT_TYPE_NOT_SUPPORTED;
        }
    } else {
        av_log(NULL, AV_LOG_ERROR, "crpt tag not found");
        return TE_ERR_DECRYPT_TAG_NOT_FOUND;
    }
}

int FFmpegXORIODecryptor::readPacket(uint8_t *buf, int bufSize) {
    if (m_pCryptIOCtx->pos + bufSize >= m_maxReadPos) {
        bufSize = static_cast<int>(m_maxReadPos - m_pCryptIOCtx->pos);
        if (bufSize <= 0) return AVERROR_EOF;
    }

    return FFmpegBaseXORIOCryptor::readPacket(buf, bufSize);
}

int64_t FFmpegXORIODecryptor::seek(int64_t offset, int whence) {
    if (SEEK_END == whence) {
        offset -= m_bdveTagSize;
    }

    auto pos = FFmpegBaseIOCryptor::seek(offset, whence);
    if (pos < 0) return pos;

    if (AVSEEK_SIZE == whence) {
        return m_maxReadPos;
    }

    return pos;
}
